package com.fanshuai.spring.boot.xa.jta;

import com.atomikos.icatch.jta.UserTransactionImp;
import com.atomikos.jdbc.AtomikosDataSourceBean;
import com.fanshuai.spring.boot.local.domain.User;
import com.fanshuai.spring.boot.sql.StatementCreator;
import lombok.extern.slf4j.Slf4j;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.stereotype.Service;
import org.springframework.transaction.TransactionDefinition;
import org.springframework.transaction.TransactionStatus;
import org.springframework.transaction.jta.JtaTransactionManager;

import javax.annotation.PostConstruct;
import javax.annotation.Resource;
import javax.transaction.UserTransaction;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.util.List;

@Service
@Slf4j
public class AtomikosXAService {
    @Resource
    private AtomikosDataSourceBean dataSourceBean1;
    @Resource
    private AtomikosDataSourceBean dataSourceBean2;

    @Resource(name = "jtaTransactionManager")
    private JtaTransactionManager jtaTransactionManager;
    @Resource
    private TransactionDefinition transactionDefinition;

    private JdbcTemplate jdbcTemplate;
    private JdbcTemplate jdbcTemplate2;

    @PostConstruct
    public void init() {
        jdbcTemplate = new JdbcTemplate(dataSourceBean1);
        jdbcTemplate2 = new JdbcTemplate(dataSourceBean2);
    }

    /**
     * jta 事务  基于JtaTransaction
     * @param userList
     */
    public void executeJtaTransactionManager(List<User> userList) {
        TransactionStatus transactionStatus = null;
        try {
            transactionStatus = jtaTransactionManager.getTransaction(transactionDefinition);
            boolean success = true;
            for (User user : userList) {
                int effectRows = jdbcTemplate.update("insert into user(name) values (?)", user.getName());
                success = success && effectRows > 0;
            }
            for (User user : userList) {
                int effectRows = jdbcTemplate2.update("insert into user(name) values (?)", user.getName());
                success = success && effectRows > 0;
            }

            if (success) {
                jtaTransactionManager.commit(transactionStatus);
            } else {
                jtaTransactionManager.rollback(transactionStatus);
            }

        } catch (Exception e) {
            if (null != transactionStatus && !transactionStatus.isCompleted()) {
                jtaTransactionManager.rollback(transactionStatus);
            }
        }
    }

    /**
     * jta 事务  基于JdbcTemplate
     * @param userList
     */
    public void executeJtaTemplate(List<User> userList) {
        UserTransaction userTransaction = null;
        try {
            userTransaction = new UserTransactionImp();
            userTransaction.begin();

            /**
             * 分布式事务分为多个本地事务，每个本地事务共享一个Connection
             */
            String sql = "insert into user(name) values (?)";
            boolean success1 = Boolean.TRUE.equals(jdbcTemplate.execute((Connection connection) -> {
                boolean b = true;
                for (User user : userList) {
                    PreparedStatement statement = connection.prepareStatement(sql);
                    statement.setString(1, user.getName());
                    int effectRows = statement.executeUpdate();
                    b = b && effectRows > 0;
                }
                return b;
            }));
            boolean success2 = Boolean.TRUE.equals(jdbcTemplate2.execute((Connection connection) -> {
                boolean b = true;
                for (User user : userList) {
                    PreparedStatement statement = connection.prepareStatement(sql);
                    statement.setString(1, user.getName());
                    int effectRows = statement.executeUpdate();
                    b = b && effectRows > 0;
                }
                return b;
            }));

            if (success1 && success2) {
                userTransaction.commit();
            } else {
                userTransaction.rollback();
            }
        } catch (Exception e) {
            e.printStackTrace();
            if (null != userTransaction) {
                try {
                    userTransaction.rollback();
                } catch (Exception ex) {
                    e.printStackTrace();
                }
            }
        }
    }

    public void executeJta(StatementCreator creator1, StatementCreator creator2) throws Exception {
        Connection connection1 = dataSourceBean1.getConnection();
        Connection connection2 = dataSourceBean2.getConnection();
        List<PreparedStatement> statements1 = creator1.createPreparedStatement(connection1);
        List<PreparedStatement> statements2 = creator2.createPreparedStatement(connection2);

        executeXA0(statements1, statements2, connection1, connection2);
    }

    private void executeXA0(List<PreparedStatement> statements1, List<PreparedStatement> statements2, Connection connection1, Connection connection2) {
        UserTransaction userTransaction = null;
        try {
            userTransaction = new UserTransactionImp();
            userTransaction.begin();

            for (PreparedStatement statement : statements1) {
                statement.execute();
            }
            for (PreparedStatement statement : statements2) {
                statement.execute();
            }

            userTransaction.commit();
        } catch (Exception e) {
            log.error("atomikos xa error, ", e);
            try {
                if (null != userTransaction) {
                    userTransaction.rollback();
                }
            } catch (Exception ex) {
                log.error("atomikos xa rollback error, ", e);
            }
        } finally {
            try {
                for (PreparedStatement statement : statements1) {
                    statement.close();
                }
                for (PreparedStatement statement : statements2) {
                    statement.close();
                }
                connection1.close();
                connection2.close();
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }
}
